#drone_model.py:
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, random_split
from collections import namedtuple
import warnings


# ====================== 可配置参数区域 (用户只需修改这里) ======================
class Config:
    # 数据配置
    FEATURE_NAMES = ['acc', 'vel', 'gyro', 'pos']  # 使用的特征
    TARGET_NAME = 'fa'  # 预测目标
    DIM_A = 3  # 特征维度

    # 训练超参数
    EPOCHS = 200  # 训练轮次
    LEARNING_RATE = 1e-4  # 学习率
    K_SHOT = 32  # 适配样本数
    ALPHA = 0.3  # 对抗损失权重
    GAMMA = 5.0  # 参数a的最大范数
    BATCH_SIZE = 16  # 训练批大小


    # 风况标签配置 (可选预定义)
    CONDITION_LABELS = {
        'no_wind': 0,
        'light_wind': 1,
        'moderate_wind': 2,
        'strong_wind': 3,
        'gale_wind': 4
    }

    # 模型保存路径
    MODEL_SAVE_PATH = './saved_models/drone_model.pth'


# ====================== 模型实现部分 (通常不需要修改) ======================
class DroneAeroModel:
    """无人机气动力与风况预测模型"""

    class PhiNet(nn.Module):
        """增强版特征提取网络"""

        def __init__(self, dim_x, dim_a):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(dim_x, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(),
                nn.Dropout(0.3),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, dim_a - 1)  # 最后一维是常数偏置
            )

        def forward(self, x):
            x = self.net(x)
            if len(x.shape) == 1:
                return torch.cat([x, torch.ones(1, device=x.device)])
            return torch.cat([x, torch.ones((x.shape[0], 1), device=x.device)], dim=-1)

    class HNet(nn.Module):
        """对抗判别网络"""

        def __init__(self, dim_a, num_c):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(dim_a, 256),
                nn.BatchNorm1d(256),
                nn.LeakyReLU(0.2),
                nn.Dropout(0.5),
                nn.Linear(256, 128),
                nn.LeakyReLU(0.2),
                nn.Linear(128, num_c)
            )

        def forward(self, x):
            return self.net(x)

    DataSample = namedtuple('DataSample', ['features', 'target', 'condition'])

    def __init__(self, config):
        """
        初始化模型
        Args:
            config: 配置类实例
        """
        self.cfg = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # 模型参数 (训练前需通过prepare_data确定)
        self.dim_x = None
        self.dim_y = None

        # 模型组件
        self.phi_net = None
        self.h_net = None
        self.is_trained = False

        # 初始化标签系统
        self._init_condition_labels()

    def _init_condition_labels(self):
        """初始化风况标签系统"""
        self.condition_map = {}
        self.inv_condition_map = {}
        self.next_condition_idx = 0

        # 加载预定义标签
        if hasattr(self.cfg, 'CONDITION_LABELS') and self.cfg.CONDITION_LABELS:
            for label, idx in self.cfg.CONDITION_LABELS.items():
                if label not in self.condition_map:
                    self.condition_map[label] = idx
                    self.inv_condition_map[idx] = label
                    if idx >= self.next_condition_idx:
                        self.next_condition_idx = idx + 1

    def add_condition_label(self, label):
        """添加风况标签"""
        if label not in self.condition_map:
            self.condition_map[label] = self.next_condition_idx
            self.inv_condition_map[self.next_condition_idx] = label
            self.next_condition_idx += 1
        return self.condition_map[label]

    def prepare_data(self, raw_data):
        """
        准备训练数据
        Args:
            raw_data: 原始数据列表，每个元素是包含特征、目标和condition字段的字典
        """
        formatted_data = []

        for data in raw_data:
            # 拼接特征
            X = np.hstack([data[feat] for feat in self.cfg.FEATURE_NAMES])
            Y = data.get(self.cfg.TARGET_NAME, np.zeros((len(X), 3)))

            # 处理风况标签
            condition = str(data.get('condition', 'unknown'))
            condition_idx = self.add_condition_label(condition)

            formatted_data.append(self.DataSample(X, Y, condition_idx))

        # 确定模型参数
        if formatted_data:
            self.dim_x = formatted_data[0].features.shape[1]
            self.dim_y = formatted_data[0].target.shape[1]

        return formatted_data

    def train(self, raw_data):
        """训练模型"""
        formatted_data = self.prepare_data(raw_data)
        if not formatted_data:
            raise ValueError("无有效训练数据")

        # 初始化模型
        self.phi_net = self.PhiNet(self.dim_x, self.cfg.DIM_A).to(self.device)
        self.h_net = self.HNet(self.cfg.DIM_A, len(self.condition_map)).to(self.device)

        # 优化器
        optimizer_phi = torch.optim.Adam(self.phi_net.parameters(), lr=self.cfg.LEARNING_RATE)
        optimizer_h = torch.optim.Adam(self.h_net.parameters(), lr=self.cfg.LEARNING_RATE)

        # 数据加载器
        train_loaders, adapt_loaders = self._create_data_loaders(formatted_data)

        # 训练循环
        for epoch in range(self.cfg.EPOCHS):
            total_loss = self._train_one_epoch(epoch, train_loaders, adapt_loaders,
                                               optimizer_phi, optimizer_h)

            # 打印进度
            if (epoch + 1) % 10 == 0:
                avg_loss_f = total_loss['f'] / len(formatted_data)
                avg_loss_c = total_loss['c'] / len(formatted_data)
                print(f"Epoch {epoch + 1}/{self.cfg.EPOCHS} - Loss_F: {avg_loss_f:.4f}, Loss_C: {avg_loss_c:.4f}")

        self.is_trained = True
        print("训练完成!")

    def _create_data_loaders(self, formatted_data):
        """创建训练和适配数据加载器"""
        train_loaders = []
        adapt_loaders = []

        for sample in formatted_data:
            dataset = self._create_dataset(sample.features, sample.target, sample.condition)

            # 划分训练集和适配集 (2:1)
            train_size = int(0.67 * len(dataset))
            train_set, adapt_set = random_split(dataset, [train_size, len(dataset) - train_size])

            train_loaders.append(DataLoader(train_set, batch_size=self.cfg.BATCH_SIZE, shuffle=True))
            adapt_loaders.append(DataLoader(adapt_set, batch_size=self.cfg.K_SHOT, shuffle=True))

        return train_loaders, adapt_loaders

    def _train_one_epoch(self, epoch, train_loaders, adapt_loaders, optimizer_phi, optimizer_h):
        """单轮训练"""
        total_loss = {'f': 0.0, 'c': 0.0}
        criterion = nn.MSELoss()
        criterion_h = nn.CrossEntropyLoss()

        # 随机遍历所有条件
        indices = np.random.permutation(len(train_loaders))

        for idx in indices:
            # 获取适配数据
            adapt_data = next(iter(adapt_loaders[idx]))
            X_adapt = adapt_data['input'].to(self.device)
            Y_adapt = adapt_data['output'].to(self.device)
            c_adapt = adapt_data['c'].to(self.device)

            # 计算适配参数a
            with torch.no_grad():
                Phi = self.phi_net(X_adapt)
                a = torch.linalg.lstsq(Phi, Y_adapt).solution
                a = self._clip_a_norm(a, self.cfg.GAMMA)

            # 获取训练数据
            train_data = next(iter(train_loaders[idx]))
            X_train = train_data['input'].to(self.device)
            Y_train = train_data['output'].to(self.device)
            c_train = train_data['c'].to(self.device)

            # 训练phi网络
            optimizer_phi.zero_grad()
            Phi_train = self.phi_net(X_train)
            Y_pred = Phi_train @ a

            loss_f = criterion(Y_pred, Y_train)
            loss_c = criterion_h(self.h_net(Phi_train), c_train)
            loss = loss_f - self.cfg.ALPHA * loss_c

            loss.backward()
            optimizer_phi.step()

            # 训练h网络 (频率较低)
            if np.random.rand() < 0.5:  # 50%概率更新
                optimizer_h.zero_grad()
                Phi_train = self.phi_net(X_train).detach()
                loss_h = criterion_h(self.h_net(Phi_train), c_train)
                loss_h.backward()
                optimizer_h.step()

            total_loss['f'] += loss_f.item()
            total_loss['c'] += loss_c.item()

        return total_loss

    def predict(self, input_data, adapt_data=None):
        """
        预测气动力和风况
        Args:
            input_data: 输入数据 [N, dim_x]
            adapt_data: 可选，适配数据 (X_adapt, Y_adapt)
        """
        if not self.is_trained:
            raise RuntimeError("模型未训练")

        # 转换输入数据
        input_tensor = torch.as_tensor(input_data, dtype=torch.float32, device=self.device)
        if input_tensor.ndim == 1:
            input_tensor = input_tensor.unsqueeze(0)

        # 计算适配参数a
        a = self._compute_adaptation_params(adapt_data) if adapt_data else \
            torch.eye(self.cfg.DIM_A, self.dim_y, device=self.device)

        # 执行预测
        with torch.no_grad():
            Phi = self.phi_net(input_tensor)
            forces = (Phi @ a).cpu().numpy()
            wind_conditions = self._predict_wind_conditions(Phi)

        return {
            'forces': forces.squeeze(),
            'wind_conditions': wind_conditions[0] if len(wind_conditions) == 1 else wind_conditions
        }

    def _compute_adaptation_params(self, adapt_data):
        """计算适配参数a"""
        X_adapt, Y_adapt = adapt_data
        X_adapt = torch.as_tensor(X_adapt, dtype=torch.float32, device=self.device)
        Y_adapt = torch.as_tensor(Y_adapt, dtype=torch.float32, device=self.device)

        with torch.no_grad():
            Phi = self.phi_net(X_adapt)
            a = torch.linalg.lstsq(Phi, Y_adapt).solution
            return self._clip_a_norm(a, self.cfg.GAMMA)

    def _predict_wind_conditions(self, Phi):
        """预测风况标签"""
        if self.h_net is None:
            return ['unknown'] * len(Phi)

        with torch.no_grad():
            wind_probs = F.softmax(self.h_net(Phi), dim=1).cpu().numpy()
            wind_labels = np.argmax(wind_probs, axis=1)
            return [self.inv_condition_map.get(idx, str(idx)) for idx in wind_labels]

    def _create_dataset(self, X, Y, C):
        """创建PyTorch数据集"""

        class AeroDataset(Dataset):
            def __init__(self, X, Y, C):
                self.X = X
                self.Y = Y
                self.C = C

            def __len__(self):
                return len(self.X)

            def __getitem__(self, idx):
                return {
                    'input': torch.FloatTensor(self.X[idx]),
                    'output': torch.FloatTensor(self.Y[idx]),
                    'c': self.C
                }

        return AeroDataset(X, Y, C)

    def _clip_a_norm(self, a, gamma):
        """限制参数a的范数"""
        norm = torch.norm(a, p='fro')
        return a if norm <= gamma else a / norm * gamma

    def save_model(self, path=None):
        """保存模型"""
        path = path or self.cfg.MODEL_SAVE_PATH
        os.makedirs(os.path.dirname(path), exist_ok=True)

        torch.save({
            'phi_state': self.phi_net.state_dict(),
            'h_state': self.h_net.state_dict(),
            'dim_x': self.dim_x,
            'dim_y': self.dim_y,
            'config': vars(self.cfg),
            'condition_map': self.condition_map,
            'next_condition_idx': self.next_condition_idx
        }, path)
        print(f"模型保存到: {path}")

    def load_model(self, path=None):
        """加载模型"""
        path = path or self.cfg.MODEL_SAVE_PATH
        checkpoint = torch.load(path, map_location=self.device)

        # 恢复配置
        for k, v in checkpoint['config'].items():
            setattr(self.cfg, k, v)

        # 恢复模型参数
        self.dim_x = checkpoint['dim_x']
        self.dim_y = checkpoint['dim_y']

        # 恢复标签系统
        self.condition_map = checkpoint['condition_map']
        self.inv_condition_map = {v: k for k, v in self.condition_map.items()}
        self.next_condition_idx = checkpoint['next_condition_idx']

        # 重建模型
        self.phi_net = self.PhiNet(self.dim_x, self.cfg.DIM_A).to(self.device)
        self.phi_net.load_state_dict(checkpoint['phi_state'])

        self.h_net = self.HNet(self.cfg.DIM_A, len(self.condition_map)).to(self.device)
        self.h_net.load_state_dict(checkpoint['h_state'])

        self.is_trained = True
        print(f"模型从 {path} 加载成功")


# ====================== 数据转换工具函数 ======================
def prepare_matlab_data(matlab_cell_data):
    """将MATLAB cell数组转换为模型需要的格式"""
    formatted_data = []

    for flight_data in matlab_cell_data:
        # 假设flight_data已经是numpy数组
        data_array = np.array(flight_data)

        sample = {
            'acc': data_array[:, 1:4],  # 加速度xyz
            'vel': data_array[:, 4:7],  # 速度xyz
            'gyro': data_array[:, 7:10],  # 陀螺仪xyz
            'pos': data_array[:, 10:13],  # 位置xyz
            'fa': data_array[:, 13:16],  # 气动力xyz(目标)
            'condition': 'light_wind'  # 设置实际风况
        }
        formatted_data.append(sample)

    return formatted_data


# ====================== 使用示例 ======================
if __name__ == "__main__":
    # 初始化配置和模型
    cfg = Config()
    model = DroneAeroModel(cfg)

    # 示例数据格式
    example_data = [
        {
            'acc': np.random.randn(100, 3),
            'vel': np.random.randn(100, 3),
            'gyro': np.random.randn(100, 3),
            'pos': np.random.randn(100, 3),
            'fa': np.random.randn(100, 3),
            'condition': 'light_wind'  # 或使用预定义的标签
        }
        # 更多数据...
    ]

    # 训练模型 (使用真实数据时取消注释)
    # model.train(example_data)
    # model.save_model()

    # 预测示例
    test_input = np.random.randn(5, 12)  # 5个样本, 12维特征
    try:
        result = model.predict(test_input)
        print("预测结果:", result)
    except RuntimeError as e:
        print(e)